#include "StdAfx.h"
#include "TelnetProtocolHandler.h"

CTelnetProtocolHandler::CTelnetProtocolHandler(void)
{
	this->allocatedPayloadCount = 1024;
	this->payloadBuffer = new unsigned char[this->allocatedPayloadCount];

	this->allocatedSubParameterCount = 1024;
	this->subParameterBuffer = new unsigned char[this->allocatedSubParameterCount];

	Reset();
}

CTelnetProtocolHandler::~CTelnetProtocolHandler(void)
{
	delete this->payloadBuffer;
	delete this->subParameterBuffer;
}

void CTelnetProtocolHandler::Reset()
{
	this->negState = NegState_Data;
	this->subState = SubState_Option;
	this->lastRequest = 0;

	this->currentSubOption = 0;
	this->subParameterCount = 0;

	this->payloadCount = 0;

	for (int i = 0; i < 256; i ++)
	{
		this->sentDX[i] = 0;
		this->sentWX[i] = 0;
		this->receivedDX[i] = 0;
		this->receivedWX[i] = 0;
	}
}

void CTelnetProtocolHandler::Transpose(unsigned char* inputBuffer, int inputCount)
{
	int outputOffset = 0;
	unsigned char* outputBuffer = new unsigned char[inputCount*2];
	VERIFY(outputBuffer != NULL);

	for (int i = 0; i < inputCount; i++)
	{
		switch (inputBuffer[i])
		{
		case IAC:
			outputBuffer[outputOffset++] = IAC;
			outputBuffer[outputOffset++] = IAC;
			break;
		case LF:	// \n -> CR LF
			if (receivedDX[TELOPT_BINARY] != DO)
			{
				outputBuffer[outputOffset++] = CR;
				outputBuffer[outputOffset++] = LF;
			} 
			else 
			{
				outputBuffer[outputOffset++] = inputBuffer[i]; // Copy verbatim in binary mode.
			}
			break;
		case CR:    // \r -> CR NUL
			if (receivedDX[TELOPT_BINARY] != DO) 
			{
				outputBuffer[outputOffset++] = CR;
				outputBuffer[outputOffset++] = NUL;
			} 
			else
			{
				outputBuffer[outputOffset++] = inputBuffer[i]; // Copy verbatim in binary mode.
			}
			break;
		default:
			outputBuffer[outputOffset++] = inputBuffer[i];
			break;
		}
	}
	Write(outputBuffer, outputOffset);
	delete outputBuffer;
}

int CTelnetProtocolHandler::Negotiate(unsigned char* inputBuffer, int inputCount)
{
	int inputOffset = 0;
	this->payloadCount = 0;

	while (inputOffset < inputCount)
	{
		unsigned char input;
		unsigned char reply;

		input = inputBuffer[inputOffset++];

		switch (this->negState)
		{
		case NegState_Data:
			if (input == IAC)
			{
				this->negState = NegState_IAC;
			}
			else 
			{
				VERIFY(IncreaseBufferIfFull(this->payloadBuffer, this->payloadCount, this->allocatedPayloadCount));
				this->payloadBuffer[this->payloadCount++] = input;
			}
			break;
		case NegState_IAC:
			switch (input)
			{
			case IAC:
				VERIFY(IncreaseBufferIfFull(this->payloadBuffer, this->payloadCount, this->allocatedPayloadCount));
				this->payloadBuffer[this->payloadCount++] = input;
				this->negState = NegState_Data;
				break;
			case WILL:
			case WONT:
			case DONT:
			case DO:
				this->lastRequest = input;
				this->negState = NegState_Negotiating;
				break;
			case EOR:
				NotifyEndOfRecord();
				this->negState = NegState_Data;
				break;
			case SB:
				this->negState = NegState_SubNegotiating;
				this->subState = SubState_Option;
				break;
			default:
				negState = NegState_Data;
				break;
			}
			break;

		case NegState_Negotiating:
			ASSERT(this->lastRequest == WILL || this->lastRequest == WONT || this->lastRequest == DO || this->lastRequest == DONT);
			switch(input)
			{
			case TELOPT_ECHO:
				SetLocalEcho(this->lastRequest == WONT || this->lastRequest == DO);
			case TELOPT_SGA:
			case TELOPT_EOR:
			case TELOPT_BINARY:
				switch (this->lastRequest)
				{
				case DO:
					reply = WILL;
					break;
				case DONT:
					reply = WONT;
					break;
				case WILL:
					reply = DO;
					break;
				case WONT:
					reply = DONT;
					break;
				default:
					ASSERT (false);
					break;
				}
				break;
			case TELOPT_NAWS:
				if (this->lastRequest == DO)
				{
					reply = WILL;
					int width, height;
					GetWindowSize(width, height);

					unsigned char replyBuffer[] = {
						IAC, reply, TELOPT_NAWS, IAC, SB, TELOPT_NAWS, 
						width >> 8, width & 0xFF, height >> 8, height & 0xFF, 
						IAC, SE
					};
					Write(replyBuffer, sizeof(replyBuffer)/sizeof(replyBuffer[0]));
				}
				break;
			default:
				reply = this->lastRequest == WILL || this->lastRequest == WONT ? DONT : WONT;
				break;
			}

			SendReplyAndUpdateRecord(this->lastRequest, reply, input);
			this->negState = NegState_Data;
			break;

		case NegState_SubNegotiating:
			switch (this->subState)
			{
			case SubState_Option:
				if (input == IAC)
				{
					this->subState = SubState_Option_IAC;
				}
				else
				{
					SubNegotiationBegin(input);
					this->subState = SubState_Data;
				}
				break;
			case SubState_Option_IAC:
				if (input == IAC)
				{
					SubNegotiationBegin(input);
					this->subState = SubState_Data;
				}
				else 
				{
					this->negState = NegState_Data;
				}
				break;
			case SubState_Data:
				if (input == IAC)
					this->subState = SubState_Data_IAC;
				else
					SubNegotiationAppendParameter(input);
				break;
			case SubState_Data_IAC:
				switch (input)
				{
				case IAC:
					this->subState = SubState_Data;
					SubNegotiationAppendParameter(input);
					break;
				case SE:
					// Handle SB
					this->negState = NegState_Data;
					break;
				case SB:
					// Handle SB
					this->subState = SubState_Option;
				default:
					break;
				}
				break;
			default:
				negState = NegState_Data;
				break;
			}
		}
	}
	return 0;
}

void CTelnetProtocolHandler::SendReplyAndUpdateRecord(unsigned char request, unsigned char reply, unsigned char option)
{
	ASSERT (request == DO || request == DONT || request == WILL || request == WONT);
	ASSERT (reply == DO || reply == DONT || reply == WILL || reply == WONT);
	ASSERT (!(request == DO || request == DONT) || (reply == WILL || reply == WONT)); // (request == DO || request == DONT) -> (reply == WILL || reply = WONT)
	ASSERT (!(reply == WILL || reply == WONT) || (request == DO || request == DONT)); // (reply == WILL || reply = WONT) -> (request == DO || request == DONT)
	ASSERT (!(request == WILL || request == WONT) || (reply == DO || reply == DONT)); // (request == WILL || request == WONT) -> (reply == DO || reply == DONT)
	ASSERT (!(reply == DO || reply == DONT) || (request == WILL || request == WONT)); // (reply == DO || reply == DONT) -> (request == WILL || request == WONT)
	
	bool needToReply = false;

	if ((request == DO || request == DONT) && (reply != this->sentWX[option] || request != this->receivedDX[option])) 
	{
		this->sentWX[option] = reply;
		this->receivedDX[option] = request;
		needToReply = true;
	}

	if ((request == WILL || request == WONT) && (reply != this->sentDX[option] || request != this->receivedWX[option]))
	{
		this->sentDX[option] = reply;
		this->receivedWX[option] = request;
		needToReply = true;
	}

	if (needToReply)
	{
		unsigned char replyBuffer[3] = {IAC, reply, option};
		Write(replyBuffer, sizeof(replyBuffer)/sizeof(replyBuffer[0]));
	}
}

void CTelnetProtocolHandler::SubNegotiationBegin(unsigned char option)
{
	this->currentSubOption = option;
	this->subParameterCount = 0;
}

void CTelnetProtocolHandler::SubNegotiationAppendParameter(unsigned char parameter)
{
	VERIFY(IncreaseBufferIfFull(this->subParameterBuffer, this->subParameterCount, this->allocatedSubParameterCount));
	this->subParameterBuffer[this->subParameterCount++] = parameter;
}

void CTelnetProtocolHandler::HandleSubNeg()
{
	switch (this->currentSubOption)
	{
	case TELOPT_TTYPE:
		if (this->subParameterCount > 0 && this->subParameterBuffer[0] == TELQUAL_SEND)
		{
			const char* terminalType = GetTerminalType();
			unsigned char buffer1[] = {IAC, SB, TELOPT_TTYPE, TELQUAL_IS};
			unsigned char buffer2[] = {IAC, SE};

			Write(buffer1, sizeof(buffer1)/sizeof(buffer1[0]));
			Write((unsigned char*)terminalType, strlen(terminalType));
			Write(buffer1, sizeof(buffer2)/sizeof(buffer2[0]));
		}
		break;
	}

	this->currentSubOption = 0;
	this->subParameterCount = 0;
}

bool CTelnetProtocolHandler::IncreaseBufferIfFull(unsigned char* &buffer, int &size, int &allocatedSize)
{
	if (size >= allocatedSize)
	{
		allocatedSize <<= 1;
		unsigned char* newBuffer = new unsigned char[allocatedSize];

		if (newBuffer == NULL)
			return false;

		memcpy(newBuffer, buffer, size);
		delete buffer;
		buffer = newBuffer;
	}

	return true;
}

void CTelnetProtocolHandler::SetLocalEcho(bool echo)
{
}

void CTelnetProtocolHandler::NotifyEndOfRecord()
{
}
